import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import sys

sys.path.append("../")
from services.DayKlineService import *
from model.StockBaseInfo import *

base = StockBaseInfo()

# 设置中文显示
plt.rcParams["font.family"] = ["SimSun"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题


def load_and_preprocess_data(code):
    """加载并预处理数据"""
    if code != None:
        try:
            klineService = DayKlineService()
            sinacode = base.translateSinaCode(code).upper()
            data = klineService.getAllDataArr(sinacode)
            if data == "" or data == None or data == False or len(data) < 1:
                return False
            stock_data = pd.DataFrame(data, columns=['id', 'code', 'Date', 'Open', 'High', 'Low', 'Close', 'Volume',
                                                     'Money', 'ma2', 'ma5', 'ma10', 'ma20', 'ma30', 'ma60', 'pre_close',
                                                     'pre_money'])
            print(stock_data)
            # 将Date列转换为datetime类型
            stock_data['Date'] = pd.to_datetime(stock_data['Date'])

            # 排序并设置日期索引
            dfall = stock_data.sort_values(by="Date")
            dfall.set_index('Date', inplace=True)
        except Exception as e:
            print(f"数据加载错误: {e}")
            # 生成示例数据用于测试
            dates = pd.date_range(start='2023-01-01', periods=365)
            values = np.random.randn(365).cumsum() + 100
            dfall = pd.DataFrame({'Close': values}, index=dates)
    else:
        # 生成示例数据
        dates = pd.date_range(start='2023-01-01', periods=365)
        values = np.random.randn(365).cumsum() + 100
        dfall = pd.DataFrame({'Close': values}, index=dates)

    return dfall


def create_dataset(data, look_back=1):
    """创建用于LSTM的数据集"""
    X, Y = [], []
    for i in range(len(data) - look_back):
        X.append(data[i:(i + look_back), 0])
        Y.append(data[i + look_back, 0])
    return np.array(X), np.array(Y)


def predict_future(model, scaler, data, look_back=30, days=5):
    """预测未来n天的股价（修改days默认值为5，look_back默认值为30）"""
    scaled_data = scaler.transform(data)
    last_sequence = scaled_data[-look_back:].reshape(1, look_back, 1)
    future_predictions = []

    for _ in range(days):  # 循环次数由days参数控制
        next_day = model.predict(last_sequence, verbose=0)
        predicted_value = next_day[0, -1, 0]
        future_predictions.append(predicted_value)
        last_sequence = np.roll(last_sequence, -1, axis=1)
        last_sequence[0, -1, 0] = predicted_value

    future_predictions = np.array(future_predictions).reshape(-1, 1)
    future_predictions = scaler.inverse_transform(future_predictions)

    return future_predictions


def plot_predictions(data, future_predictions, look_back=30):
    """可视化历史和预测股价，包括历史预测线"""
    plt.figure(figsize=(14, 7))

    # 绘制实际历史股价
    plt.plot(data.index, data['Close'], label='实际股价', color='blue')

    # 准备历史预测数据
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data[['Close']])

    # 创建历史预测数据
    X, _ = create_dataset(scaled_data, look_back)

    # 检查是否有足够的数据进行历史预测
    if len(X) == 0:
        print("警告: 数据长度不足以生成历史预测")
    else:
        # 重塑数据以适应LSTM模型
        X = np.reshape(X, (X.shape[0], X.shape[1], 1))

        # 加载模型以进行历史预测
        model = tf.keras.models.load_model('stock_prediction_model.keras')

        # 预测历史数据（输出为3维数组）
        historical_predictions = model.predict(X, verbose=0)

        # 关键修改：提取每个样本的最后一个时间步，转换为2维数组
        historical_predictions = historical_predictions[:, -1, 0].reshape(-1, 1)
        historical_predictions = scaler.inverse_transform(historical_predictions)

        # 创建历史预测的日期索引
        hist_dates = data.index[look_back:]

        # 绘制历史预测线
        plt.plot(hist_dates, historical_predictions, label='历史预测', color='green', alpha=0.7)

    # 绘制未来预测线
    last_date = data.index[-1]
    future_dates = pd.date_range(start=last_date + timedelta(days=1), periods=len(future_predictions))
    plt.plot(future_dates, future_predictions, 'r--o', label='未来预测')

    # 连接最后一个实际点和第一个未来预测点
    if len(data) > 0 and len(future_predictions) > 0:
        plt.plot([data.index[-1], future_dates[0]], [data['Close'].iloc[-1], future_predictions[0, 0]], 'k--',
                 alpha=0.5)

    plt.title('股价预测与历史表现')
    plt.xlabel('日期')
    plt.ylabel('价格')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def main():
    # 加载模型
    model = tf.keras.models.load_model('stock_prediction_model.keras')
    print("模型已加载")

    # 加载数据
    data = load_and_preprocess_data(sys.argv[1])

    # 保存原始数据用于绘图
    original_data = data.copy()

    # 数据标准化
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data[['Close']])

    # 预测未来5天的股价（修改days参数为5，look_back为30）
    look_back = 30
    future_predictions = predict_future(model, scaler, data[['Close']], look_back, days=5)

    # 打印未来5天的预测股价
    print("未来5天的预测股价:")
    last_date = data.index[-1]
    for i, price in enumerate(future_predictions, 1):
        future_date = last_date + timedelta(days=i)
        print(f"{future_date.date()}: {price[0]:.2f}")

    # 可视化预测结果
    plot_predictions(original_data, future_predictions, look_back)


if __name__ == "__main__":
    main()